"""
训练ResNet网络
"""
from torchvision import transforms
from torch import nn, cuda, optim, no_grad, load
from torch.nn import functional as F
# 自定义
from model.ResNet import resnet18, resnet34, resnet50, resnet101, resnet152
from model.MobileNetv2 import mobilenetv2
from model.MobileNetv3 import mobilenetv3
from PIL import Image
from argparse import ArgumentParser
from collections import OrderedDict
import os


# 加载训练集与测试集
def predict(args):
    #------------------------------------------------------------#
    # GPU is available?
    #------------------------------------------------------------#
    flag = True if cuda.is_available() else False

    #------------------------------------------------------------#
    # data augment
    #------------------------------------------------------------#
    compose = transforms.Compose([
            transforms.Resize((args.size, args.size)),
            transforms.ToTensor()
        ])
    # ------------------------------------------------------------#
    # init model
    # ------------------------------------------------------------#
    model = eval(args.model)(args.classes)
    state = load(os.path.join(os.getcwd(), args.weights))
    new_state = OrderedDict()
    for k, m in state.items():
        new_state[k[7:]] = m
    model.load_state_dict(new_state)
    if flag:
        model.cuda()

    with no_grad():
        # 声明非训练模式
        model.eval()
        # 打开图片
        image = Image.open(os.path.join(os.getcwd(), args.source))
        trans_image = compose(image)
        trans_image = trans_image.cuda()
        forward_ret = model(trans_image.unsqueeze(0))
        out = F.softmax(forward_ret, 1)
        index = out.argmax(1)
        print("classes is {}".format(index))

        # 计算正确个数

if __name__ =="__main__":
    # 初始化参数解析器
    parser = ArgumentParser(description="training config")

    parser.add_argument("--model", type=str, default="resnet18", help="model")
    parser.add_argument("--classes", type=int, default=400, help="classes number")
    parser.add_argument("--weights", type=str, default="./runs/resnet18/weights/best.pth", help="weights")
    parser.add_argument("--size", type=int, default=1140, help="image size")
    parser.add_argument("--source", type=str, default="./experiments/04x06.jpg", help="save path")

    # 解析参数
    args = parser.parse_args()
    predict(args)